﻿from memory import GlobalMemory
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
from enum import Enum
import numpy as np
from numpy import ndarray, uint8, uint16, uint32, uint64, float32, float64
import threading
from datetime import datetime
from typedef import *


class PMETA(Enum):
    # 页表meta的枚举
    Dirty = 1
    Loaded = 2
    CacheEN = 4
    Accessed = 8


class MMU_Stat(Enum):
    IDLE = 0
    L1_4k = 1
    L2_4k = 2
    L1_2m = 3
    L2_2m = 4
    L1_1g = 5
    L2_1g = 6
    MEM_GET = 7
    DONE = 8


@dataclass
class TLBEntry:
    ppfn: int = 0
    access: int = 0
    meta: PMETA = 0

    def __init__(self, value: int):
        self.access = (value >> 4) & 0x0f
        self.meta = (value >> 8) & 0x0f
        self.ppfn = (value >> 12) & 0x0fff

    def __int__(self):
        return (self.ppfn << 12) | (self.meta << 8) | (self.access << 4)


@dataclass
class LRU_TLB:
    def __init__(self, size: int, memory: GlobalMemory):
        self.__mem = memory
        self.__size = size
        self.__data: Dict[int, List[TLBEntry, int]] = {}
        self.__history: Dict[int, int] = {}

    def __getitem__(self, tag: int) -> List[TLBEntry, int]:
        if tag not in self.__data.keys():
            if len(self.__data.keys()) == self.__size:
                id = min(self.__history.keys(), key=self.__history.get)
                self.__history.pop(id)
                self.__data.pop(id)
            self.__data[tag] = [TLBEntry(), 0]
        self.__history[tag] = datetime.now().timestamp()
        return self.__data[tag]

    def __setitem__(self, tag: int, value: List[TLBEntry, int]):
        if tag in self.__data.keys():
            if self.__data[tag][0].meta & PMETA.Dirty:
                self.__mem.write(int(self.__data[tag][0]), self.__data[tag][1])
        self.__data[tag] = value
        self.__history[tag] = datetime.now().timestamp()

    def delete(self):
        for i in self.__data.keys():
            if self.__data[i][0].meta & PMETA.Dirty:
                self.__mem.write(int(self.__data[i][0]), self.__data[i][1])
        self.__data.clear()
        self.__history.clear()


@dataclass
class MMU:
    def __init__(self, memory: GlobalMemory, reglist: List[int] = None):
        self.__mem = memory
        self.__l1i: Dict[str, Dict[int, LRU_TLB]] = {
            '4k': {}, '2m': {}}  # 16it*16way/it*2
        self.__l1d: Dict[str, Dict[int, LRU_TLB]] = {
            '4k': {}, '2m': {}, '1g': {}}  # (16it+8it+8it)*16way/it
        self.__l2: Dict[str, Dict[int, LRU_TLB]] = {
            '4k': {}, '2m': {}, '1g': {}}  # (128it+64it+64it)*8way/it
        self.__regs: List[int] = reglist

    def update(self, reglist: List[int]):
        self.__regs = reglist

    def __call__(self, addr: int, access: int) -> int:
        if self.__regs[Status] ^ PageEnable:  # 页表未启用
            return addr & 0xf_ffff
        pfn = addr >> 12
        stat: MMU_Stat = MMU_Stat.IDLE
        idx = 0
        while stat != MMU_Stat.DONE:
            match stat:
                case MMU_Stat.IDLE:
                    if addr < self.__regs[RvBA]:
                        raise Exception(
                            f"Access denied {access} at {addr}:Reversed Virtual Pages")
                    stat = MMU_Stat.L1_4k
                case MMU_Stat.L1_4k:
                    idx = pfn & 0x0f
                    if self.__l1i['4k'].get(idx) is None:
                        self.__l1i['4k'][idx] = LRU_TLB(16, self.__mem)
                    if self.__l1d['4k'].get(idx) is None:
                        self.__l1d['4k'][idx] = LRU_TLB(16, self.__mem)
                    line = self.__l1i['4k'][idx][pfn]\
                        if access & Execute\
                        else self.__l1d['4k'][idx][pfn]
                    if line[0].meta & PMETA.Loaded:  # 页表已找到
                        if (line[0].access & access) != access:
                            raise AccessViolate(addr)
                        if access & Write:
                            line[0].meta |= PMETA.Dirty
                        line[0].meta |= PMETA.Accessed
                        return (line[0].ppfn << 12) | (addr & 0xfff)
                    else:
                        stat = MMU_Stat.L2_4k
                case MMU_Stat.L2_4k:
                    idx = pfn & 0x7f
                    line = self.__l2['4k'][idx][pfn]
                    if line[0].meta & PMETA.Loaded:  # 页表已找到
                        if (line[0].access & access) != access:
                            raise AccessViolate(addr)
                        if access & Write:
                            line[0].meta |= PMETA.Dirty
                        line[0].meta |= PMETA.Accessed
                        if access & Execute:
                            self.__l1i['4k'][pfn & 0x0f][pfn] = line
                        else:
                            self.__l1d['4k'][pfn & 0x0f][pfn] = line
                        return (line[0].ppfn << 12) | (addr & 0xfff)
                    else:
                        stat = MMU_Stat.L1_2m
                case MMU_Stat.L1_2m:
                    line = self.__l1i['2m'][pfn & 0x0f][pfn]\
                        if access & Execute\
                        else self.__l1d['2m'][pfn & 0x07][pfn]
                    if line[0].meta & PMETA.Loaded:  # 页表已找到
                        if (line[0].access & access) != access:
                            raise AccessViolate(addr)
                        if access & Write:
                            line[0].meta |= PMETA.Dirty
                        line[0].meta |= PMETA.Accessed
                        return (line[0].ppfn << 12) | (addr & 0xfff)
                    else:
                        stat = MMU_Stat.L2_2m
                case MMU_Stat.L2_2m:
                    idx = pfn & 0x3f
                    line = self.__l2['2m'][idx][pfn]
                    if line[0].meta & PMETA.Loaded:  # 页表已找到
                        if (line[0].access & access) != access:
                            raise AccessViolate(addr)
                        if access & Write:
                            line[0].meta |= PMETA.Dirty
                        line[0].meta |= PMETA.Accessed
                        if access & Execute:
                            self.__l1i['2m'][pfn & 0x0f][pfn] = line
                        else:
                            self.__l1d['2m'][pfn & 0x0f][pfn] = line
                        return (line[0].ppfn << 12) | (addr & 0xfff)
                    else:
                        stat = MMU_Stat.L1_1g
                case MMU_Stat.L1_1g:
                    if access & Execute:
                        stat = MMU_Stat.L2_1g
                    else:
                        idx = pfn & 0x07
                        line = self.__l1d['1g'][idx][pfn]
                        if line[0].meta & PMETA.Loaded:  # 页表已找到
                            if (line[0].access & access) != access:
                                raise AccessViolate(addr)
                            if access & Write:
                                line[0].meta |= PMETA.Dirty
                            line[0].meta |= PMETA.Accessed
                            return (line[0].ppfn << 12) | (addr & 0xfff)
                        else:
                            stat = MMU_Stat.L2_1g
                case MMU_Stat.L2_1g:
                    idx = pfn & 0x3f
                    line = self.__l2['1g'][idx][pfn]
                    if line[0].meta & PMETA.Loaded:  # 页表已找到
                        if (line[0].access & access) != access:
                            raise AccessViolate(addr)
                        if access & Write:
                            line[0].meta |= PMETA.Dirty
                        line[0].meta |= PMETA.Accessed
                        if access & Execute:
                            pass
                        else:
                            self.__l1d['1g'][pfn & 0x0f][pfn] = line
                        return (line[0].ppfn << 12) | (addr & 0xfff)
                    else:
                        stat = MMU_Stat.MEM_GET
                case MMU_Stat.MEM_GET:  # 从内存获取页表
                    level: int = 3 if (pfn < 0x200) \
                        else 2 if (pfn < 0x40000) \
                        else 1 if (pfn < 0x800_0000) \
                        else 0
                    page_idx = self.__regs[PTR0+level]
                    temp: TLBEntry = TLBEntry(0)
                    while level < 4:
                        temp = TLBEntry(self.__mem.read(page_idx, 'qword'))
                        if temp.access & (Read | Write | Execute):  # 页表有效
                            if temp.meta & PMETA.Loaded:  # 页面已经加载
                                if (temp.access & access) != access:
                                    raise AccessViolate(addr)
                                if access & Write:
                                    temp.meta |= PMETA.Dirty
                                temp.meta |= PMETA.Accessed
                                match level:
                                    case 1: self.__l2['1g'][pfn & 0x3f][pfn] = [temp, page_idx]
                                    case 2: self.__l2['2m'][pfn & 0x3f][pfn] = [temp, page_idx]
                                    case 3: self.__l2['4k'][pfn & 0x3f][pfn] = [temp, page_idx]
                                return (temp.ppfn << 12) | (addr & 0xfff)
                            else:
                                raise PageFault(addr)
                    raise IllegalPage(addr)
                case MMU_Stat.DONE:
                    return addr
